//! 注释字符串的迭代器
//! struct AnnotatedString 将为我们保存一个字符串和一些注释
//! 我们希望遍历它以返回 AnnotatedStringPart

use unicode_segmentation::UnicodeSegmentation;

use crate::{
    editor::highlighter::{AnnotationType, annotation::AnnotatedStr},
    prelude::ByteIdx,
};

#[derive(Debug)]
pub struct AnnotatedStrPart<'a> {
    pub string: &'a str,
    pub annotation_type: Option<AnnotationType>,
}

pub struct AnnotatedStrIter<'a> {
    // 'a: 对 AnnotatedString 的引用生命周期应至少与迭代器本身一样长。
    pub annotated_string: &'a AnnotatedStr,
    pub current_idx: ByteIdx,
}

impl<'a> Iterator for AnnotatedStrIter<'a> {
    type Item = AnnotatedStrPart<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current_idx >= self.annotated_string.string.len() {
            return None;
        }

        // 获取所有字素的起始字节索引
        let grapheme_indices: Vec<usize> = self
            .annotated_string
            .string
            .grapheme_indices(true)
            .map(|(idx, _)| idx)
            .collect();

        // 找到当前索引所在的字素位置
        let current_pos = grapheme_indices
            .iter()
            .position(|&idx| idx >= self.current_idx)
            .unwrap_or(grapheme_indices.len());

        let start_idx = grapheme_indices
            .get(current_pos)
            .copied()
            .unwrap_or(self.current_idx);

        // 检查当前位置是否有注释
        if let Some(annotation) = self
            .annotated_string
            .annotations
            .iter()
            .filter(|a| a.start <= self.current_idx && a.end > self.current_idx)
            .last()
        {
            // 找到注释结束位置对应的字素边界
            let end_pos = grapheme_indices
                .iter()
                .position(|&idx| idx >= annotation.end)
                .unwrap_or(grapheme_indices.len());
            let end_idx = grapheme_indices
                .get(end_pos)
                .copied()
                .unwrap_or(self.annotated_string.string.len());

            self.current_idx = end_idx;
            return Some(AnnotatedStrPart {
                string: &self.annotated_string.string[start_idx..end_idx],
                annotation_type: Some(annotation.annotation_type),
            });
        }

        // 没有注释时，找到下一个注释的起始位置或字符串末尾
        let mut end_idx = self.annotated_string.string.len();
        for annotation in &self.annotated_string.annotations {
            if annotation.start > self.current_idx && annotation.start < end_idx {
                end_idx = annotation.start;
            }
        }

        // 确保结束索引在字素边界上
        let end_pos = grapheme_indices
            .iter()
            .position(|&idx| idx >= end_idx)
            .unwrap_or(grapheme_indices.len());
        let end_idx = grapheme_indices.get(end_pos).copied().unwrap_or(end_idx);

        self.current_idx = end_idx;
        Some(AnnotatedStrPart {
            string: &self.annotated_string.string[start_idx..end_idx],
            annotation_type: None,
        })
    }
}
